# -*- coding:utf-8 -*-
import sys
import os
import sqlite3
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *


class DataModel(QAbstractTableModel):
    def __init__(self, data, header):
        super(DataModel, self).__init__()

        if data is not None:
            self.row_count = len(data)
            self.col_count = len(data[0])
        else:
            self.row_count = 0
            self.col_count = 0

        self.data = data
        self.header = header

    def rowCount(self, QModelIndex_parent=None, *args, **kwargs):
        return self.row_count

    def columnCount(self, QModelIndex_parent=None, *args, **kwargs):
        return self.col_count

    def headerData(self, p_int, Qt_Orientation, int_role=None):
        if int_role == Qt.DisplayRole and Qt_Orientation == Qt.Horizontal:
            if self.header is None:
                return ""
            return self.header[p_int][0]
        else:
            return None

    def data(self, QModelIndex, int_role=None):
        if int_role == Qt.DisplayRole:
            try:
                if self.data is None:
                    return ""
                return self.data[QModelIndex.row()][QModelIndex.column()]
            except:
                return None
        else:
            return None


class TreeItem(object):
    def __init__(self, data, parent=None):
        self.parentItem = parent
        self.itemData = data
        self.childItems = []

    def appendChild(self, item):
        self.childItems.append(item)

    def child(self, row):
        return self.childItems[row]

    def childCount(self):
        return len(self.childItems)

    def columnCount(self):
        return len(self.itemData)

    def data(self, column):
        try:
            return self.itemData[column]
        except IndexError:
            return None

    def parent(self):
        return self.parentItem

    def row(self):
        if self.parentItem:
            return self.parentItem.childItems.index(self)

        return 0


class TreeModel(QAbstractItemModel):
    def __init__(self, data, parent=None):
        super(TreeModel, self).__init__(parent)

        #self.rootItem = TreeItem(("Title", "Summary"))
        self.rootItem = TreeItem(("Navigator", ))
        #self.setupModelData(data.split('\n'), self.rootItem)
        if data == None:
            lines = [("No table", ), ]
        else:
            lines = data
        self.setupModelData(lines, self.rootItem)

    def columnCount(self, parent):
        if parent.isValid():
            return parent.internalPointer().columnCount()
        else:
            return self.rootItem.columnCount()

    def data(self, index, role):
        if not index.isValid():
            return None

        if role != Qt.DisplayRole:
            return None

        item = index.internalPointer()

        return item.data(index.column())

    def flags(self, index):
        if not index.isValid():
            return Qt.NoItemFlags

        return Qt.ItemIsEnabled | Qt.ItemIsSelectable

    def headerData(self, section, orientation, role):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return self.rootItem.data(section)

        return None

    def index(self, row, column, parent):
        if not self.hasIndex(row, column, parent):
            return QModelIndex()

        if not parent.isValid():
            parentItem = self.rootItem
        else:
            parentItem = parent.internalPointer()

        childItem = parentItem.child(row)
        if childItem:
            return self.createIndex(row, column, childItem)
        else:
            return QModelIndex()

    def parent(self, index):
        if not index.isValid():
            return QModelIndex()

        childItem = index.internalPointer()
        parentItem = childItem.parent()

        if parentItem == self.rootItem:
            return QModelIndex()

        return self.createIndex(parentItem.row(), 0, parentItem)

    def rowCount(self, parent):
        if parent.column() > 0:
            return 0

        if not parent.isValid():
            parentItem = self.rootItem
        else:
            parentItem = parent.internalPointer()

        return parentItem.childCount()

    def setupModelData(self, lines, parent):
        parents = [parent]

        for line in lines:
            parents[-1].appendChild(TreeItem(line, parents[-1]))


class Form(QWidget):

    executeCompleted = pyqtSignal(str)

    def __init__(self, parent=None):
        super(Form, self).__init__(parent)

        self.connection = None

        self.sqline = QLineEdit()
        self.sqline.setDisabled(True)
        self.sqline.returnPressed.connect(self.execute)

        self.executeButton = QPushButton("Execute")
        self.executeButton.setDisabled(True)
        self.executeButton.clicked.connect(self.execute)

        execLayout = QHBoxLayout()
        execLayout.addWidget(self.sqline)
        execLayout.addWidget(self.executeButton)

        self.grid = QTableView()
        self.tree = QTreeView()
        self.tree.setMaximumWidth(200)

        resultLayout = QVBoxLayout()
        resultLayout.addLayout(execLayout)
        resultLayout.addWidget(self.grid)

        viewLayout = QHBoxLayout()
        viewLayout.addWidget(self.tree)
        viewLayout.addLayout(resultLayout)

        mainLayout = QHBoxLayout()
        mainLayout.addLayout(execLayout)
        mainLayout.addLayout(viewLayout)
        self.setLayout(mainLayout)
        self.setOpened(False)

    @pyqtSlot(bool)
    def setOpened(self, enabled):
        self.sqline.setEnabled(enabled)
        self.executeButton.setEnabled(enabled)
        self.grid.setEnabled(enabled)
        self.setEnabled(enabled)
        if enabled is True:
            model = TreeModel(self._getTableList())
            self.tree.setModel(model)

    def _getTableList(self):
        sql = "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
        cursor = self.connection.execute(sql)
        table = []
        for row in cursor.fetchall():
            table.append(row)
        cursor.close()
        return table

    def execute(self):
        sql = self.sqline.text()
        msg = ""
        if sql == "":
            return

        if self.connection is None:
            QMessageBox.information(self, "No database", "Open or create a database")
            return

        try:
            cursor = self.connection.cursor()
            cursor.execute(sql)
            if cursor.description is not None: # data query
                # TODO: Binding data grid
                self.grid.setModel(DataModel(cursor.fetchall(), cursor.description))
                msg = "Query succeed"
            else:
                self.grid.setModel(DataModel(None, None))
                msg = "The current sql has executed successfully"
            cursor.close()
            self.connection.commit()
        except:
            msg = "Execute SQL error"
        self.executeCompleted.emit(msg)


class MainWindow(QMainWindow):
    databaseOpened = pyqtSignal(bool)
    def __init__(self):
        super(MainWindow, self).__init__()

        self.openAction = QAction("Open", self)
        self.openAction.setShortcut('Ctrl+O')
        self.openAction.setStatusTip('Open a database')
        self.openAction.triggered.connect(self.openDatabase)

        self.createAction = QAction("Create", self)
        self.createAction.setShortcut('Ctrl+R')
        self.createAction.setStatusTip("Create a database")
        self.createAction.triggered.connect(self.createDatabase)

        self.closeAction = QAction("Close", self)
        self.closeAction.setShortcut('Ctrl+Q')
        self.closeAction.setStatusTip("Close a database")
        self.closeAction.triggered.connect(self.closeDatabase)
        self.closeAction.setEnabled(False)

        menu = self.menuBar()
        fileMenu = menu.addMenu("&File")
        fileMenu.addAction(self.openAction)
        fileMenu.addAction(self.createAction)
        fileMenu.addAction(self.closeAction)

        aboutAction = QAction('&About', self)
        aboutAction.triggered.connect(self.about)
        menu.addAction(aboutAction)

        self.form = Form()
        self.form.executeCompleted.connect(self.setStatus)
        self.databaseOpened.connect(self.form.setOpened)
        self.setCentralWidget(self.form)
        self.setWindowTitle("PySQLite")
        self.showMaximized()

    def openDatabase(self):
        fileName, _ = QFileDialog.getOpenFileName(self, "Open Database", "./", "SQLite database (*.db)")
        self.form.connection = sqlite3.connect(fileName)
        if self.form.connection is not None:
            self.setStatus("Succeed open database at %s" % fileName)
            self.openAction.setEnabled(False)
            self.createAction.setEnabled(False)
            self.closeAction.setEnabled(True)
            self.databaseOpened.emit(True)

    def closeDatabase(self):
        if self.form.connection is not None:
            self.form.connection.close()
            self.openAction.setEnabled(True)
            self.createAction.setEnabled(True)
            self.closeAction.setEnabled(False)
            self.databaseOpened.emit(False)

    def createDatabase(self):
        fileName, _ = QFileDialog.getSaveFileName(self, "Create Database", "./", "SQLite database (*.db)")
        self.form.connection = sqlite3.connect(fileName)
        self.setStatus("Succeed create database at %s" % fileName)
        self.openAction.setEnabled(False)
        self.createAction.setEnabled(False)
        self.closeAction.setEnabled(True)
        self.databaseOpened.emit(True)

    @pyqtSlot(str)
    def setStatus(self, msg):
        status = self.statusBar()
        status.showMessage(msg)

    def about(self):
        QMessageBox.information(self, "About PySQLite", "Author: Kelvin Fang")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    sys.exit(app.exec_())
